from typing import List,Dict, Any, Optional,Union
from neo4j import AsyncDriver
from core.config import settings

class Neo4jService:
    """Neo4j图数据库服务类"""

    def __init__(self,driver:AsyncDriver):
        self.driver = driver

    async def execute_query(
            self,
            query: str,
            parameters: Optional[Dict] = None,
            database: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """执行Cypher查询并返回结果"""
        try:
            # 使用异步会话
            async with self.driver.session(
                    database=database or settings.NEO4J_DATABASE
            ) as session:
                result = await session.run(query, parameters or {})
                records = await result.data()
                return records

        except Exception as e:
            print(f"Neo4j 查询执行失败: {str(e)}")
            raise

    # async def keyword_search(
    #         self,
    #         keyword: str,
    #         database: Optional[str] = None
    # ) -> List[Dict[str, Any]]:
    #     """
    #     使用关键词进行搜索（简化版）
    #
    #     参数:
    #     - keyword: 搜索关键词
    #     - database: 数据库名称 (可选)
    #
    #     返回:
    #     - 匹配结果的字典列表
    #     """
    #     # 使用您提供的查询模板
    #     query = """
    #     MATCH (a:xianxiang)-[r1]->(b)
    #     WHERE a.xxmc CONTAINS $kw OR b.xxmc CONTAINS $kw OR b.gzmc CONTAINS $kw
    #     RETURN a.xxmc AS head, type(r1) AS rel, coalesce(b.xxmc, '') + coalesce(b.gzmc, '') AS tail
    #     UNION
    #     MATCH (a:guzhang)-[r2]->(b)
    #     WHERE a.gzmc CONTAINS $kw OR b.wtmc CONTAINS $kw
    #     RETURN a.gzmc AS head, type(r2) AS rel, coalesce(b.wtmc, '') AS tail
    #     UNION
    #     MATCH (a:wenti)-[r3]->(b)
    #     WHERE a.wtmc CONTAINS $kw
    #     RETURN a.wtmc AS head, type(r3) AS rel, '' AS tail
    #     UNION
    #     MATCH (a:xianxiangfenzhi)-[r4]->(b)
    #     WHERE a.xxfzmc CONTAINS $kw OR b.gzmc CONTAINS $kw OR b.wtmc CONTAINS $kw
    #     RETURN a.xxfzmc AS head, type(r4) AS rel, coalesce(b.gzmc, '') + coalesce(b.wtmc, '') AS tail
    #     """
    #
    #     try:
    #         # 执行查询
    #         results = await self.execute_query(
    #             query,
    #             parameters={"kw": keyword},
    #             database=database
    #         )
    #
    #         # 格式化结果
    #         triples = []
    #         for record in results:
    #             head = record["head"]
    #             rel = record["rel"]
    #             tail = record["tail"]
    #
    #             if head and tail:
    #                 triples.append({
    #                     "head": head,
    #                     "relation": rel,
    #                     "tail": tail
    #                 })
    #
    #         return triples
    #
    #     except Exception as e:
    #         print(f"关键词搜索失败: {str(e)}")
    #         raise

    async def keyword_search(
            self,
            keyword: str,
            database: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        使用关键词进行搜索，返回D3兼容的图格式

        参数:
        - keyword: 搜索关键词
        - database: 数据库名称 (可选)

        返回:
        - 包含"nodes"和"links"的字典，适合前端图谱可视化
        """
        # 使用您提供的查询模板
        query = """
        MATCH (a:xianxiang)-[r1]->(b)
        WHERE a.xxmc CONTAINS $kw OR b.xxmc CONTAINS $kw OR b.gzmc CONTAINS $kw
        RETURN a.xxmc AS head, type(r1) AS rel, coalesce(b.xxmc, '') + coalesce(b.gzmc, '') AS tail, 
               labels(a) as head_labels, labels(b) as tail_labels
        UNION
        MATCH (a:guzhang)-[r2]->(b)
        WHERE a.gzmc CONTAINS $kw OR b.wtmc CONTAINS $kw
        RETURN a.gzmc AS head, type(r2) AS rel, coalesce(b.wtmc, '') AS tail,
               labels(a) as head_labels, labels(b) as tail_labels
        UNION
        MATCH (a:wenti)-[r3]->(b)
        WHERE a.wtmc CONTAINS $kw
        RETURN a.wtmc AS head, type(r3) AS rel, '' AS tail,
               labels(a) as head_labels, labels(b) as tail_labels
        UNION
        MATCH (a:xianxiangfenzhi)-[r4]->(b)
        WHERE a.xxfzmc CONTAINS $kw OR b.gzmc CONTAINS $kw OR b.wtmc CONTAINS $kw
        RETURN a.xxfzmc AS head, type(r4) AS rel, coalesce(b.gzmc, '') + coalesce(b.wtmc, '') AS tail,
               labels(a) as head_labels, labels(b) as tail_labels
        """

        try:
            # 执行查询
            results = await self.execute_query(
                query,
                parameters={"kw": keyword},
                database=database
            )

            # 用于存储去重后的节点和关系
            nodes_dict = {}  # key: 节点ID, value: 节点数据
            links = []

            for record in results:
                head = record["head"]
                rel = record["rel"]
                tail = record["tail"]
                head_labels = record.get("head_labels", [])
                tail_labels = record.get("tail_labels", [])

                # 只有head和tail都存在时才处理
                if head and tail:
                    # 处理头节点
                    if head not in nodes_dict:
                        node_type = self._infer_node_type(head, head_labels)
                        nodes_dict[head] = {
                            "id": head,
                            "name": head,
                            "type": node_type,
                            "labels": head_labels
                        }

                    # 处理尾节点
                    if tail not in nodes_dict:
                        node_type = self._infer_node_type(tail, tail_labels)
                        nodes_dict[tail] = {
                            "id": tail,
                            "name": tail,
                            "type": node_type,
                            "labels": tail_labels
                        }

                    # 创建关系链接
                    links.append({
                        "source": head,
                        "target": tail,
                        "relation": rel,
                        "type": "relationship"
                    })

            # 转换字典为列表
            nodes = list(nodes_dict.values())

            return {
                "nodes": nodes,
                "links": links
            }

        except Exception as e:
            print(f"关键词搜索失败: {str(e)}")
            raise

    def _infer_node_type(self, node_name: str, labels: List[str]) -> str:
        """
        根据节点名称和标签推断节点类型
        """
        # 基于标签推断
        if labels:
            label_str = "".join(labels).lower()
            if "xianxiang" in label_str:
                return "现象"
            elif "guzhang" in label_str:
                return "故障"
            elif "wenti" in label_str:
                return "问题"
            elif "xianxiangfenzhi" in label_str:
                return "现象分支"

        # 基于关键词推断（备用方案）
        node_lower = node_name.lower()
        if any(keyword in node_lower for keyword in ["飞轮", "转速", "燃烧"]):
            return "现象"
        elif any(keyword in node_lower for keyword in ["故障", "损坏", "异常"]):
            return "故障"
        elif any(keyword in node_lower for keyword in ["问题", "原因", "因素"]):
            return "问题"
        else:
            return "实体"  # 默认类型



